# langchain框架中的链条---（Chains） 可以理解为对输入的prompt在框架的层面进行if-elif-else的动作。再优化prompt，再输入给LLM，理论上可以通过多个LLM进行回答。

import pandas as pd
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain
from langchain.chains import SimpleSequentialChain
from langchain.chains import SequentialChain
from langchain.chains.router import MultiPromptChain
from langchain.chains.router.llm_router import LLMRouterChain,RouterOutputParser
from langchain.prompts import PromptTemplate
import warnings
warnings.filterwarnings('ignore')

api_key = "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"
df = pd.read_csv('Data.csv')

# LLM链（LLMChain）
def get_langchain_chain_1():
    chat = ChatOpenAI(api_key=api_key, temperature=0.9)
    prompt = ChatPromptTemplate.from_template(
        "What is the best name to describe \
        a company that makes {product}?"
    )
    chain = LLMChain(llm=chat, prompt=prompt)
    product = "Queen Size Sheet Set"
    response = chain.run(product)
    print(response)

# 简单顺序链（SimpleSequentialChain）
def get_langchain_chain_2():
    chat = ChatOpenAI(api_key=api_key, temperature=0.9)
    # prompt template 1
    first_prompt = ChatPromptTemplate.from_template(
        "What is the best name to describe \
        a company that makes {product}?"
    )
    # Chain 1
    chain_one = LLMChain(llm=chat, prompt=first_prompt)
    # prompt template 2
    second_prompt = ChatPromptTemplate.from_template(
        "Write a 20 words description for the following \
        company:{company_name}"
    )
    # chain 2
    chain_two = LLMChain(llm=chat, prompt=second_prompt)
    overall_simple_chain = SimpleSequentialChain(chains=[chain_one, chain_two],
                                                 verbose=True
                                                 )
    product = "Queen Size Sheet Set"
    response = overall_simple_chain.run(product)
    print(response)

# 常规顺序链（SequentialChain）
def get_langchain_chain_3():
    chat = ChatOpenAI(api_key=api_key, temperature=0.9)

    # prompt template 1: translate to english
    first_prompt = ChatPromptTemplate.from_template(
        "Translate the following review to english:"
        "\n\n{Review}"
    )
    # chain 1: input= Review and output= English_Review
    chain_one = LLMChain(llm=chat, prompt=first_prompt,
                         output_key="English_Review"
                         )
    second_prompt = ChatPromptTemplate.from_template(
        "Can you summarize the following review in 1 sentence:"
        "\n\n{English_Review}"
    )
    # chain 2: input= English_Review and output= summary
    chain_two = LLMChain(llm=chat, prompt=second_prompt,
                         output_key="summary"
                         )
    # prompt template 3: translate to english
    third_prompt = ChatPromptTemplate.from_template(
        "What language is the following review:\n\n{Review}"
    )
    # chain 3: input= Review and output= language
    chain_three = LLMChain(llm=chat, prompt=third_prompt,
                           output_key="language"
                           )
    # prompt template 4: follow up message
    fourth_prompt = ChatPromptTemplate.from_template(
        "Write a follow up response to the following "
        "summary in the specified language:"
        "\n\nSummary: {summary}\n\nLanguage: {language}"
    )
    # chain 4: input= summary, language and output= followup_message
    chain_four = LLMChain(llm=chat, prompt=fourth_prompt,
                          output_key="followup_message"
                          )
    # overall_chain: input= Review
    # and output= English_Review,summary, followup_message
    overall_chain = SequentialChain(
        chains=[chain_one, chain_two, chain_three, chain_four],
        input_variables=["Review"],
        output_variables=["English_Review", "summary", "followup_message"],
        verbose=True
    )
    review = df.Review[5]
    overall_chain(review)


# 路由链（Router Chain）
def get_langchain_chain_4():
    physics_template = """You are a very smart physics professor. \
    You are great at answering questions about physics in a concise\
    and easy to understand manner. \
    When you don't know the answer to a question you admit\
    that you don't know.

    Here is a question:
    {input}"""

    math_template = """You are a very good mathematician. \
    You are great at answering math questions. \
    You are so good because you are able to break down \
    hard problems into their component parts, 
    answer the component parts, and then put them together\
    to answer the broader question.

    Here is a question:
    {input}"""

    history_template = """You are a very good historian. \
    You have an excellent knowledge of and understanding of people,\
    events and contexts from a range of historical periods. \
    You have the ability to think, reflect, debate, discuss and \
    evaluate the past. You have a respect for historical evidence\
    and the ability to make use of it to support your explanations \
    and judgements.

    Here is a question:
    {input}"""

    computerscience_template = """ You are a successful computer scientist.\
    You have a passion for creativity, collaboration,\
    forward-thinking, confidence, strong problem-solving capabilities,\
    understanding of theories and algorithms, and excellent communication \
    skills. You are great at answering coding questions. \
    You are so good because you know how to solve a problem by \
    describing the solution in imperative steps \
    that a machine can easily interpret and you know how to \
    choose a solution that has a good balance between \
    time complexity and space complexity. 

    Here is a question:
    {input}"""

    prompt_infos = [
        {
            "name": "physics",
            "description": "Good for answering questions about physics",
            "prompt_template": physics_template
        },
        {
            "name": "math",
            "description": "Good for answering math questions",
            "prompt_template": math_template
        },
        {
            "name": "History",
            "description": "Good for answering history questions",
            "prompt_template": history_template
        },
        {
            "name": "computer science",
            "description": "Good for answering computer science questions",
            "prompt_template": computerscience_template
        }
    ]

    chat = ChatOpenAI(api_key=api_key, temperature=0)

    destination_chains = {}
    for p_info in prompt_infos:
        name = p_info["name"]
        prompt_template = p_info["prompt_template"]
        prompt = ChatPromptTemplate.from_template(template=prompt_template)
        chain = LLMChain(llm=chat, prompt=prompt)
        destination_chains[name] = chain

    destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]
    destinations_str = "\n".join(destinations)

    default_prompt = ChatPromptTemplate.from_template("{input}")
    default_chain = LLMChain(llm=chat, prompt=default_prompt)

    MULTI_PROMPT_ROUTER_TEMPLATE = """Given a raw text input to a \
    language model select the model prompt best suited for the input. \
    You will be given the names of the available prompts and a \
    description of what the prompt is best suited for. \
    You may also revise the original input if you think that revising\
    it will ultimately lead to a better response from the language model.

    << FORMATTING >>
    Return a markdown code snippet with a JSON object formatted to look like:
    ```json
    {{{{
        "destination": string \ name of the prompt to use or "DEFAULT"
        "next_inputs": string \ a potentially modified version of the original input
    }}}}
    ```

    REMEMBER: "destination" MUST be one of the candidate prompt \
    names specified below OR it can be "DEFAULT" if the input is not\
    well suited for any of the candidate prompts.
    REMEMBER: "next_inputs" can just be the original input \
    if you don't think any modifications are needed.

    << CANDIDATE PROMPTS >>
    {destinations}

    << INPUT >>
    {{input}}

    << OUTPUT (remember to include the ```json)>>"""

    router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(
        destinations=destinations_str
    )
    router_prompt = PromptTemplate(
        template=router_template,
        input_variables=["input"],
        output_parser=RouterOutputParser(),
    )
    router_chain = LLMRouterChain.from_llm(chat, router_prompt)

    chain = MultiPromptChain(router_chain=router_chain,
                             destination_chains=destination_chains,
                             default_chain=default_chain, verbose=True
                             )

    # response = chain.run("what is you name")
    response = chain.run("What is black body radiation?")
    # response = chain.run("what is 2 + 2")
    # response = chain.run("Why does every cell in our body contain DNA?")
    print(response)

if __name__ == '__main__':
    # print(df.head())
    # get_langchain_chain_1()
    # get_langchain_chain_2()
    # get_langchain_chain_3()
    get_langchain_chain_4()